/-
Copyright (c) 2022 Jannis Limperg. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Jannis Limperg
-/

import Aesop.Tree.Data

open Lean
open Lean.Meta

namespace Aesop

/-
The following functions let us run MetaM actions in the context of a rapp or
goal. Rapps save the metavariable context in which they were run by storing a
`Meta.SavedState`. When we, for example, apply a rule to a goal, we run the
rule's action in the metavariable context of the goal (which is the
metavariable context of the goal's parent rapp). The resulting metavariable
context, in which the goal mvar is assigned to an expression generated by the
rule, then becomes the metavariable context of the rule's rapp.

To save and restore metavariable contexts, we use the `MonadBacktrack MetaM`
instance. This means that some elements of the state are persistent, notably
caches and trace messages. These become part of the global state.

The environment is not persistent. This means that modifications of the
environment made by a rule are not visible in the global state and are reset
once the tactic exits. As a result, rules which modify the environment are
likely to fail.
-/

def Rapp.runMetaM (x : MetaM α) (r : Rapp) : MetaM (α × Meta.SavedState) :=
  r.metaState.runMetaM x

def Rapp.runMetaM' (x : MetaM α) (r : Rapp) : MetaM α :=
  Prod.fst <$> r.runMetaM x

def Rapp.runMetaMModifying (x : MetaM α) (r : Rapp) : MetaM (α × Rapp) := do
  let (result, finalState) ← r.runMetaM x
  return (result, r |>.setMetaState finalState)

def RappRef.runMetaMModifying (x : MetaM α) (rref : RappRef) : MetaM α := do
  let (result, r) ← (← rref.get).runMetaMModifying x
  rref.set r
  return result

def Goal.runMetaMInPostNormState (x : MVarId → MetaM α) (g : Goal) :
    MetaM (α × Meta.SavedState) := do
  let some (postGoal, postState) := g.postNormGoalAndMetaState? | throwError
    "aesop: internal error: expected goal {g.id} to be normalised (but not proven by normalisation)."
  postState.runMetaM $ x postGoal

def Goal.runMetaMInPostNormState' (x : MVarId → MetaM α) (g : Goal) : MetaM α :=
  Prod.fst <$> g.runMetaMInPostNormState x

def Goal.runMetaMInParentState (x : MetaM α) (g : Goal) :
    MetaM (α × Meta.SavedState) := do
  match ← g.parentRapp? with
  | none => withoutModifyingState' x
  | some rref => (← rref.get).runMetaM x

def Goal.runMetaMInParentState' (x : MetaM α) (g : Goal) :
    MetaM α :=
  Prod.fst <$> g.runMetaMInParentState x

def Goal.runMetaMModifyingParentState (x : MetaM α) (g : Goal) :
    MetaM α := do
  match ← g.parentRapp? with
  | none => x
  | some rref => rref.runMetaMModifying x

def Rapp.runMetaMInParentState (x : MetaM α) (r : Rapp) :
    MetaM (α × Meta.SavedState) := do
  (← r.parent.get).runMetaMInParentState x

def Rapp.runMetaMInParentState' (x : MetaM α) (r : Rapp) :
    MetaM α := do
  (← r.parent.get).runMetaMInParentState' x

def Rapp.runMetaMModifyingParentState (x : MetaM α) (r : Rapp) :
    MetaM α := do
  (← r.parent.get).runMetaMModifyingParentState x

end Aesop
